"""
处理数据的一些相关操作
"""

import os
import json
import numpy as np
from tqdm import tqdm

import torch
import torchaudio

import hyperparameters


def casia_wav_route():
    """ 统计 CASIA数据集 所有 wav 文件的路径信息
        并将路径信息存为 json 文件

        统计信息：4名录制者，每位录制者每种情绪录制 50 个音频文件
    """

    # 如果 json 文件已经存在则跳过
    json_name = os.path.join(hyperparameters.BASE_DIRECTORY, "casia_wav_route.json")
    if not os.path.exists(json_name):  # 文件存在则不执行
        
        wav_route = {}

        dir = os.path.join(hyperparameters.BASE_DIRECTORY, "casia")

        for subject in os.listdir(dir):  # 录音者姓名
            if not subject.endswith(".txt"):  # 排除掉唯一一个 txt 文件
                wav_route[subject] = {}
                wav_path = os.path.join(dir, subject)
                for label in os.listdir(wav_path):
                    if not label.endswith(".ini"):  # 排除掉唯一一个 ini 文件
                        wav_route[subject][label] = []
                        file_path = os.path.join(wav_path, label)
                        for filename in os.listdir(file_path):
                            if filename.endswith(".wav"):
                                wav_route[subject][label].append(os.path.join(file_path, filename))

        # 将统计完的路径信息写入 data 目录下的 json 文件
        with open(json_name, "w") as f:
            f.write(json.dumps(wav_route, indent=4))


def read_json(directory):
    """ 读取 json 文件的函数
    
    inputs:
        directory: json 文件的目录
    output:
        content: json 文件的内容
    """
    
    with open(directory) as f:
        content = json.load(f)
    return content


def normalize_wav_data(wave_data, sample_rate):
    """ 规范化 wav 数据，切割或填充为相同的长度大小

    inputs:
        wave_data: Tensor 数据，形状为 (channel, time)
        samplt_rate: 音频采样率
    """

    # 目标音频长度
    real_len = wave_data.shape[1]
    target_len = hyperparameters.SEGMENT_NUM * sample_rate

    if real_len == target_len:  # 长度刚好
        return wave_data
    elif real_len < target_len:  # 长度不够，两端平均填充
        pre = int((target_len - real_len) / 2)
        pos = target_len - real_len - pre
        return torch.cat([torch.zeros(1, pre), wave_data, torch.zeros(1, pos)], dim=1)
    else:  # 长度过长，取中间数据
        pre = int((real_len - target_len) / 2)
        pos = real_len - target_len - pre
        return wave_data.split([pre, target_len, pos], dim=1)[1]


def write_wave(wave_data, wave_path, sample_rate):
    """  将划分好的数据转回 wav 音频文件存到所属类别的文件夹下

    inputs:
        wave_data: 划分好的数据
        wave_path: 数据所存的路径
        sample_rate: 采样率
    """
    torchaudio.save(wave_path, wave_data, sample_rate)


def segment_dataset(wav_dict):
    """ 对数据集进行训练集与测试集的切分
        将切分规范好的音频文件放到对应的文件夹下

    inputs:
        wav_dict: 存有 wav 音频文件的字典
    """

    segment_path = os.path.join(
        hyperparameters.BASE_DIRECTORY, "{}_{}_segment".format(
            hyperparameters.DATA_NAME, hyperparameters.SEGMENT_NUM
        )
    )

    # 如果数据集未完成切分则进行切分
    if not os.path.exists(segment_path):
        # 创建文件夹
        os.makedirs(segment_path)
        os.makedirs(os.path.join(segment_path, "train"))
        os.makedirs(os.path.join(segment_path, "test"))
        for label in hyperparameters.CASIA_LABELS:
            os.makedirs(os.path.join(segment_path, "train", label))
            os.makedirs(os.path.join(segment_path, "test", label))

        # 开始便利文件
        for name, labels in tqdm(wav_dict.items()):  # 加一个处理进度条
            for label, files in labels.items():
                np.random.shuffle(files)  # 将原始的路径打乱，保证测试集选择的随机性
                boundary = int(len(files) * hyperparameters.TRAIN_RATIO)
                
                # 训练集数据写入
                for file in files[:boundary]:
                    wave_data, sample_rate = torchaudio.load(file)
                    # 规范长度后的数据
                    wave_data = normalize_wav_data(wave_data, sample_rate)
                    # 规范后的数据写入对应路径下
                    write_wave(
                        wave_data,
                        os.path.join(segment_path, "train", label, f"{name}_{os.path.basename(file)}"),
                        sample_rate
                    )
                
                # 测试集数据写入
                for file in files[boundary:]:
                    wave_data, sample_rate = torchaudio.load(file)
                    wave_data = normalize_wav_data(wave_data, sample_rate)
                    # 规范后的数据写入对应路径下
                    write_wave(
                        wave_data,
                        os.path.join(segment_path, "test", label, f"{name}_{os.path.basename(file)}"),
                        sample_rate
                    )


if __name__ == "__main__":
    casia_wav_route()

    casia_wav = read_json("data/casia_wav_route.json")

    segment_dataset(casia_wav)
